From 348d723d8b61d1b5a40b44aa607c8898efbd27b5 Mon Sep 17 00:00:00 2001
From: Chazeon <chenxing.luo@gmail.com>
Date: Tue, 5 Jan 2021 01:09:05 -0500
Subject: [PATCH] Support data extraction at geotherm PT.

---
 README.md           | 18 +++++++++--------
 cij/cli/cij.py      |  3 +++
 cij/cli/extract.py  |  2 +-
 cij/cli/geotherm.py | 47 +++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 61 insertions(+), 9 deletions(-)
 create mode 100644 cij/cli/geotherm.py

diff --git a/README.md b/README.md
index f41d33e..b3bf55c 100644
--- a/README.md
+++ b/README.md
@@ -21,25 +21,27 @@ Options:
   --help     Show this message and exit.
 
 Commands:
-  extract     Extract the value and create table for multiple variables at...
-  fill        Fill non-zero Cij terms based on symmetry.
-  modes       Plot interpolated mode frequency vs volume.
-  plot        Plot SAM-Cij calculation results.
-  run         Perform SAM-Cij calculation.
-  run-static  Calculate elastic moduli and acoustic velocities.
+  extract           Create data table at specific P or T.
+  extract-geotherm  Create data table at geotherm PT.
+  fill              Fill non-zero Cij terms based on symmetry.
+  modes             Plot interpolated mode frequency vs volume.
+  plot              Plot SAM-Cij calculation results.
+  run               Perform SAM-Cij calculation.
+  run-static        Calculate elastic moduli and acoustic velocities.
 ```
 
 And are avaliable as standalone commands:
 
 - **`cij-run`**         - Perform SAM-Cij calculation.
 - **`cij-run-static`**  - Calculate elastic moduli and acoustic velocities.
-- **`cij-extract`**     - Extract the value and create table for multiple variables at spcific P or T.
+- **`cij-extract`**     - Extract the value and create data table for multiple variables at geotherm PT.
+- **`cij-extract-geotherm`** - Extract the value and create data table for multiple variables at spcific P or T.
 - **`cij-fill`**        - Fill non-zero Cij terms based on symmetry.
 - **`cij-modes`**       - Plot interpolated mode frequency vs volume.
 - **`cij-plot`**        - Plot SAM-Cij calculation results.
 
 
-### SAM-Cij calculations with `qha-cij` or `cij run`
+### SAM-Cij calculations with `cij run`
 
 #### Configuration file
 
diff --git a/cij/cli/cij.py b/cij/cli/cij.py
index b86f04d..a5ccf8c 100644
--- a/cij/cli/cij.py
+++ b/cij/cli/cij.py
@@ -20,6 +20,9 @@ def main():
 from cij.cli.extract import main as _extract
 main.add_command(_extract, "extract")
 
+from cij.cli.geotherm import main as _geotherm
+main.add_command(_geotherm, "extract-geotherm")
+
 from cij.cli.modes import main as _modes
 main.add_command(_modes, "modes")
 
diff --git a/cij/cli/extract.py b/cij/cli/extract.py
index 3be9cdf..171b41e 100644
--- a/cij/cli/extract.py
+++ b/cij/cli/extract.py
@@ -12,7 +12,7 @@ def load_data(var):
     df.index = [float(idx) for idx in df.index]
     return df
 
-@click.command(help="Extract the value and create table for multiple variables at specific P or T.")
+@click.command(help="Create data table at specific P or T.")
 @click.option("-T", "--temperature", type=click.FLOAT, help="Specify temperature in K.")
 @click.option("-P", "--pressure", type=click.FLOAT, help="Specify tressure in GPa.")
 @click.option("-v", "--variables", required=True, help="Variables to output.")
diff --git a/cij/cli/geotherm.py b/cij/cli/geotherm.py
new file mode 100644
index 0000000..46fb8bb
--- /dev/null
+++ b/cij/cli/geotherm.py
@@ -0,0 +1,47 @@
+import click
+from typing import List
+
+
+def load_data(var):
+
+    import pandas
+    from glob import glob
+
+    df = pandas.read_table(glob(f"{var}_tp_*")[0], sep=r"\s+", index_col=0)
+    df.columns = [float(colname) for colname in df.columns]
+    df.index = [float(idx) for idx in df.index]
+    return df
+
+def fit_data(df):
+
+    from scipy.interpolate import RectBivariateSpline
+
+    x = df.index.to_numpy()
+    y = df.columns.to_numpy()
+    z = df.to_numpy()
+    return RectBivariateSpline(x, y, z)
+
+@click.command(help="Create data table at geotherm PT.")
+@click.option("-i", "-g", "--geotherm", required=True, help="The file name of geotherm P, D, T map.", type=click.Path(exists=True))
+@click.option("--t-col", help="The name of geotherm pressure column.", default="P")
+@click.option("--p-col", help="The name of geotherm temperature column", default="T")
+@click.option("-v", "--variables", required=True, help="Variables to output.")
+@click.option("-h", "--hide-header", default=False, is_flag=True, help="Hide header or not.")
+def main(variables: List[str], hide_header: bool, t_col: str = None, p_col: str = None, geotherm: str = None):
+
+    import glob
+    import pandas
+    import numpy
+
+
+    variables = variables.split(",")
+    table = pandas.read_table(geotherm, sep=r"\s+", index_col=None, header=0)
+
+    for var in variables:
+        df = load_data(var)
+        table[var] = fit_data(df)(table[p_col], table[t_col], grid=False)
+
+    print(table.to_string(header=(not hide_header), index=False))
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file