update polars (#211)

Makefile

@@ -12,10 +12,10 @@ venv:
 clean:
 	@rm -r venv
 
-test:
+test: venv
 	$(PYTHON_BIN)/pytest pygef/*
 	$(PYTHON_BIN)/pytest pygef/robertson/*
 	$(PYTHON_BIN)/pytest pygef/been_jefferies/*
 
-pre-commit:
-	$(python) black .
+pre-commit: venv
+	$(PYTHON_BIN)/black .

pygef/expressions.py

@@ -24,11 +24,10 @@ def ic_to_gamma(water_level):
     )
 
 
-def ic_to_soil_type():
+def ic_to_soil_type(ti: pl.Expr = pl.col("type_index")):
     """
     Assign the soil type to the corresponding Ic.
     """
-    ti = pl.col("type_index")
     return (
         pl.when(ti > 3.22)
         .then("Peat")

pygef/plot_utils.py

@@ -1,29 +1,35 @@
+# needed for python 3.7
+from __future__ import annotations
 from copy import copy
 
 import matplotlib.patches as mpatches
 import matplotlib.pyplot as plt
 import numpy as np
 import polars as pl
 
-colours_robertson = {
-    "Peat": "#a76b29",
-    "Clays - silty clay to clay": "#578E57",
-    "Silt mixtures - clayey silt to silty clay": "#0078C1",
-    "Sand mixtures - silty sand to sandy silt": "#DBAD4B",
-    "Sands - clean sand to silty sand": "#e5c581",
-    "Gravelly sand to dense sand": "#708090",
-    None: "black",
-}
-
-colours_been_jefferies = {
-    "Peat": "#a76b29",
-    "Clays": "#578E57",
-    "Clayey silt to silty clay": "#0078C1",
-    "Silty sand to sandy silt": "#DBAD4B",
-    "Sands: clean sand to silty": "#e5c581",
-    "Gravelly sands": "#708090",
-    None: "black",
-}
+
+def colors_robertson() -> dict[str, str]:
+    return {
+        "Peat": "#a76b29",
+        "Clays - silty clay to clay": "#578E57",
+        "Silt mixtures - clayey silt to silty clay": "#0078C1",
+        "Sand mixtures - silty sand to sandy silt": "#DBAD4B",
+        "Sands - clean sand to silty sand": "#e5c581",
+        "Gravelly sand to dense sand": "#708090",
+        None: "black",
+    }
+
+
+def colors_been_jefferies() -> dict[str, str]:
+    return {
+        "Peat": "#a76b29",
+        "Clays": "#578E57",
+        "Clayey silt to silty clay": "#0078C1",
+        "Silty sand to sandy silt": "#DBAD4B",
+        "Sands: clean sand to silty": "#e5c581",
+        "Gravelly sands": "#708090",
+        None: "black",
+    }
 
 
 def num_columns(classification, df_group):
@@ -58,6 +64,7 @@ def plot_cpt(
     title_group = None
 
     if classification is not None:
+        print(df)
         df, title = assign_color(df, classification, colors)
     if df_group is not None:
         df_group = copy(df_group)
@@ -142,21 +149,19 @@ def assign_color(df, classification, colors):
     """
     if colors is None:
         if classification == "robertson":
+            colors = colors_robertson()
             df = df.with_column(
-                pl.col("soil_type")
-                .apply(lambda row: colours_robertson[row])
-                .alias("colour")
+                pl.col("soil_type").apply(lambda row: colors[row]).alias("colour")
             )
 
             return (
                 df,
                 "Robertson",
             )
         elif classification == "been_jefferies":
+            colors = colors_been_jefferies()
             df = df.with_column(
-                pl.col("soil_type")
-                .apply(lambda row: colours_been_jefferies[row])
-                .alias("colour")
+                pl.col("soil_type").apply(lambda row: colors[row]).alias("colour")
             )
 
             return (
@@ -184,7 +189,7 @@ def add_plot_classification(fig, df, depth_max, depth_min, title, num_col, z_NAP
     :return:
     """
     ax = fig.add_subplot(1, num_col, 3)
-    df[df["soil_type"].is_null(), "soil_type"] = "UNKNOWN"
+    df = df.with_column(pl.col("soil_type").fill_null("UNKNOWN"))
     for st in df["soil_type"].unique():
         partial_df = df[df["soil_type"] == st]
         if z_NAP:
@@ -333,10 +338,10 @@ def plot_bore(df, figsize=(11, 8), show=True, dpi=100):
     return fig
 
 
-def get_legend(classification, colors):
+def get_legend(classification, colors) -> dict[str, str]:
     if colors is not None:
         return colors
     elif classification == "robertson":
-        return colours_robertson
+        return colors_robertson()
     elif classification == "been_jefferies":
-        return colours_been_jefferies
+        return colors_been_jefferies()

pygef/robertson/tests.py

@@ -92,8 +92,9 @@ def test_ic_to_gamma(self):
             assert v[column] == df[column]
 
     def test_ic_to_soil_type(self):
-        df1 = pl.DataFrame({"type_index": [2.208177, 2.408926, 2.793642]})
-        v = util.ic_to_soil_type(df1)
+        v = pl.DataFrame({"type_index": [2.208177, 2.408926, 2.793642]}).with_column(
+            util.ic_to_soil_type()
+        )
         df = pl.DataFrame(
             {
                 "type_index": [2.208177, 2.408926, 2.793642],

pygef/robertson/util.py

@@ -38,35 +38,27 @@ def type_index(df):
     return df
 
 
-def ic_to_soil_type(df):
+def ic_to_soil_type(ti: pl.Expr = pl.col("type_index")):
     """
     Assign the soil type to the corresponding Ic.
-
-    :param df: (DataFrame) Original DataFrame.
-    :return: (DataFrame) Updated DataFrame.
     """
-    # TODO: how to fill it properly with the same initial values?
-    df["soil_type"] = np.tile("", len(df.rows()))
-
-    ic_mask = df["type_index"] > 3.6
-    df[ic_mask, "soil_type"] = "Peat"
-
-    ic_mask = df["type_index"] <= 3.6
-    df[ic_mask, "soil_type"] = "Clays - silty clay to clay"
-
-    ic_mask = df["type_index"] <= 2.95
-    df[ic_mask, "soil_type"] = "Silt mixtures - clayey silt to silty clay"
-
-    ic_mask = df["type_index"] <= 2.6
-    df[ic_mask, "soil_type"] = "Sand mixtures - silty sand to sandy silt"
 
-    ic_mask = df["type_index"] <= 2.05
-    df[ic_mask, "soil_type"] = "Sands - clean sand to silty sand"
-
-    ic_mask = df["type_index"] <= 1.31
-    df[ic_mask, "soil_type"] = "Gravelly sand to dense sand"
-
-    return df
+    return (
+        pl.when(ti > 3.6)
+        .then("Peat")
+        .when((ti <= 3.6) & (ti > 2.95))
+        .then("Clays - silty clay to clay")
+        .when((ti <= 2.95) & (ti > 2.6))
+        .then("Silt mixtures - clayey silt to silty clay")
+        .when((ti <= 2.6) & (ti > 2.05))
+        .then("Sand mixtures - silty sand to sandy silt")
+        .when((ti <= 2.05) & (ti > 1.31))
+        .then("Sands - clean sand to silty sand")
+        .when(ti <= 1.31)
+        .then("Gravelly sand to dense sand")
+        .otherwise("")
+        .alias("soil_type")
+    )
 
 
 def none_to_zero(df):
@@ -138,7 +130,7 @@ def condition(x):
             if new:
                 n = df["n"]
 
-    return df.pipe(ic_to_soil_type)
+    return df.with_column(ic_to_soil_type())
 
 
 def old_robertson(

requirements.txt

@@ -1,6 +1,6 @@
 matplotlib >= 3.4.2
 pandas >= 1.1.5
-polars[pyarrow] >= 0.13.31
+polars[pyarrow] >= 0.13.55
 xmlschema == 1.11.2
 Sphinx >= 3.3.1
 sphinx-autodoc-typehints >= 1.11.1
@@ -9,3 +9,4 @@ asteroid-sphinx-theme >= 0.0.3
 sphinx_rtd_theme >= 1.0.0
 pytest>=6.2.5
 lxml>=4.7.1
+black==22.6.0

setup.py

@@ -11,7 +11,7 @@
     license="mit",
     packages=["pygef", "pygef.been_jefferies", "pygef.robertson"],
     install_requires=[
-        "polars>= 0.9.5",
+        "polars>= 0.13.55",
         "matplotlib>= 3.4.2",
         "lxml==4.9.1",
     ],